#include "Graphics/sdlgpu.h"
#include "Graphics/renderer.h"
#include "SDL3_image/SDL_image.h"
#include "glm/glm.hpp"
#include "glm/gtc/matrix_transform.hpp"

using namespace Jin;

struct Vertex {
	glm::vec3 pos;
	glm::vec2 uv;
};

struct MVP {
	glm::mat4 model;
	glm::mat4 view;
	glm::mat4 proj;
};

struct Plane {
	glm::vec3 position;
	glm::vec3 rotation;
	glm::vec3 scale = glm::vec3(1, 1, 1);
	glm::vec4 color;
	std::shared_ptr<GPUTexture> texture;
};

struct FlyCamera {
	void MoveTo(const glm::vec3& p) {
		position = p;
	}

	void Move(const glm::vec3& offset) {
		position += offset;
	}

	void RotateX(float angle) {
		rotation.x += angle;
		rotation.x = glm::clamp(rotation.x, -89.0f, 89.0f);
	}

	void RotateY(float angle) {
		rotation.y += angle;
	}

	const glm::mat4& GetMat() const {
		return mat;
	}

	void Update() {
		mat =
			glm::translate(
				glm::rotate(
					glm::rotate(
						glm::rotate(glm::mat4(1.0), -glm::radians(rotation.x), glm::vec3(1, 0, 0)),
						-glm::radians(rotation.y), glm::vec3(0, 1, 0)),
					-glm::radians(rotation.z), glm::vec3(0, 0, 1)),
				-position);
	}

	glm::mat4 mat = glm::mat4(1.0);
	glm::vec3 position = glm::vec3(0, 0, 0);
	glm::vec3 rotation = glm::vec3(0, 0, 0);
};

struct AppContext {

	SDLFramework sdlfx;
	std::unique_ptr<SDLWindow> window;
	std::unique_ptr<GPUDevice> gpuDevice;
	std::unique_ptr<GPUBuffer> vertexBuffer;
	std::shared_ptr<GPUTexture> windowTexture;
	std::shared_ptr<GPUTexture> floorTexture;
	std::unique_ptr<GPUTexture> depthTexture;

	MVP mvp{};
	float rotateY = 0.0f;
	float rotateX = 0.0f;

	std::vector<Plane> planes;
	FlyCamera camera;

	AppContext() {
		printf("SDL version: %d.%d.%d\n", sdlfx.version.major, sdlfx.version.minor, sdlfx.version.micor);
	}

	bool Init() {
		sdlfx.SetFPS(200);

		window = std::make_unique<SDLWindow>(_T("testSDLGPU"), 800, 600, 0);
		if (!*window) {
			return false;
		}

		sdlfx.onFPS = [this](Uint32 fps) {
			window->SetTitle(std::format(_T("fps:{}, delta:{}"), fps, sdlfx.deltaTime));
			};

		if (!InitGPUDevice()) {
			return false;
		}
		UploadVertexData();
		if (!vertexBuffer || !*vertexBuffer) {
			return false;
		}
		windowTexture = CreateImageTexture("images/window.png");
		if (!windowTexture || !*windowTexture) {
			return false;
		}
		floorTexture = CreateImageTexture("images/floor.png");
		if (!floorTexture || !*floorTexture) {
			return false;
		}
		CreateDepthTexture(800, 600);
		if (!depthTexture || !*depthTexture) {
			return false;
		}

		mvp.model = glm::mat4(1.0f);
		mvp.view = glm::mat4(1.0f);
		//mvp.view = glm::lookAt(glm::vec3(0.0f, 0.0f, 0.0f), glm::vec3(0.0f, 0.0f, -1.0f), glm::vec3(0.0f, 1.0f, 0.0f));
		mvp.proj = glm::perspective(glm::radians(45.0f), 800.0f / 600.0f, 0.1f, 100.0f);

		InitPlanes();

		return true;
	}

	bool InitGPUDevice() {
		gpuDevice = std::make_unique<GPUDevice>(*window);
		if (!*gpuDevice) {
			printf("Create GPU Device failed!\n");
			return false;
		}
		printf("GPU Device Driver: %s\n", gpuDevice->GetDeviceDriver());

		gpuDevice->SetSwapchainParameters(SDL_GPU_SWAPCHAINCOMPOSITION_SDR, SDL_GPU_PRESENTMODE_IMMEDIATE);

		gpuDevice->onLoadShaderCode = [](std::vector<Uint8>& vsCode, std::vector<Uint8>& fsCode, stringv driver) {
			SDLFileStorage storage("shaders");
			if (!storage) {
				printf("open shaders storage failed! %s\n", SDL_GetError());
				return false;
			}
			if (!storage.ReadFile("shader5.vert.spv", vsCode)) {
				printf("read shader5.vert.spv failed! %s\n", SDL_GetError());
				return false;
			}
			if (!storage.ReadFile("shader5.frag.spv", fsCode)) {
				printf("read shader5.frag.spv failed! %s\n", SDL_GetError());
				return false;
			}
			return true;
			};
		gpuDevice->onSetShaderCreateInfo = [](GPUShaderCreateInfo& vsInfo, GPUShaderCreateInfo& fsInfo, stringv driver) {
			vsInfo.num_uniform_buffers = 1;
			fsInfo.num_uniform_buffers = 1;
			fsInfo.num_samplers = 1;
			};
		if (!gpuDevice->CreateShaders()) {
			return false;
		}

		gpuDevice->onSetVertexBufferDescriptions = [](GPUVertexBufferDescription*& descs, Uint32& num) {
			descs[0].slot = 0;
			descs[0].input_rate = SDL_GPU_VERTEXINPUTRATE_VERTEX;
			descs[0].instance_step_rate = 0;
			descs[0].pitch = sizeof(Vertex);
			num = 1;
			};
		gpuDevice->onSetVertexAttributes = [](GPUVertexAttribute*& attrs, Uint32& num) {
			attrs[0].location = 0;
			attrs[0].buffer_slot = 0;
			attrs[0].format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT3;
			attrs[0].offset = offsetof(Vertex, pos);
			attrs[1].location = 1;
			attrs[1].buffer_slot = 0;
			attrs[1].format = SDL_GPU_VERTEXELEMENTFORMAT_FLOAT2;
			attrs[1].offset = offsetof(Vertex, uv);
			num = 2;
			};
		gpuDevice->onSetRasterizerState = [](GPURasterizerState& state) {
			state.cull_mode = SDL_GPU_CULLMODE_BACK;
			state.fill_mode = SDL_GPU_FILLMODE_FILL;
			state.front_face = SDL_GPU_FRONTFACE_COUNTER_CLOCKWISE;
			};
		gpuDevice->onSetDepthStencilState = [](GPUDepthStencilState& state) {
			state.compare_op = SDL_GPU_COMPAREOP_LESS;
			state.back_stencil_state.fail_op = SDL_GPU_STENCILOP_ZERO;
			state.back_stencil_state.pass_op = SDL_GPU_STENCILOP_ZERO;
			state.back_stencil_state.depth_fail_op = SDL_GPU_STENCILOP_ZERO;
			state.back_stencil_state.compare_op = SDL_GPU_COMPAREOP_NEVER;
			state.compare_mask = 0xff;
			state.write_mask = 0xff;
			state.enable_depth_test = true;
			state.enable_depth_write = true;
			state.enable_stencil_test = false;
			};
		gpuDevice->onSetColorTargetDescriptions = [](GPUColorTargetDescription*& descs, Uint32& num) {
			descs[0].blend_state.src_color_blendfactor = SDL_GPU_BLENDFACTOR_ONE;
			descs[0].blend_state.dst_color_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA;
			descs[0].blend_state.src_alpha_blendfactor = SDL_GPU_BLENDFACTOR_ONE;
			descs[0].blend_state.dst_alpha_blendfactor = SDL_GPU_BLENDFACTOR_ZERO;
			num = 1;
			};
		gpuDevice->onSetGraphicsPipelineTargetInfo = [](GPUGraphicsPipelineTargetInfo& targetInfo) {
			targetInfo.has_depth_stencil_target = true;
			targetInfo.depth_stencil_format = SDL_GPU_TEXTUREFORMAT_D16_UNORM;
			};
		if (!gpuDevice->CreateGraphicsPipeline()) {
			return false;
		}
		if (!gpuDevice->CreateSampler()) {
			return false;
		}
		return true;
	}

	void UploadVertexData() {
		Vertex vertices[]{
			Vertex{ { -0.5f, -0.5f, 0.0f }, { 0.0f, 0.0f} },
			Vertex{ {  0.5f, -0.5f, 0.0f }, { 1.0f, 0.0f} },
			Vertex{ {  0.5f,  0.5f, 0.0f }, { 1.0f, 1.0f} },
			Vertex{ {  0.5f,  0.5f, 0.0f }, { 1.0f, 1.0f} },
			Vertex{ { -0.5f,  0.5f, 0.0f }, { 0.0f, 1.0f} },
			Vertex{ { -0.5f, -0.5f, 0.0f }, { 0.0f, 0.0f} },
		};

		GPUTransferBuffer transferBuffer{ *gpuDevice, SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD, sizeof(vertices) };
		if (!transferBuffer) {
			printf("Create transfer buffer failed! %s\n", SDL_GetError());
			return;
		}
		memcpy(GPUTransferBufferMapper{ transferBuffer }, vertices, sizeof(vertices));

		vertexBuffer = std::make_unique<GPUBuffer>(*gpuDevice, SDL_GPU_BUFFERUSAGE_VERTEX, (Uint32)sizeof(vertices));
		if (!*vertexBuffer) {
			printf("Create vertex buffer failed! %s\n", SDL_GetError());
			return;
		}

		GPUCommandBuffer cmd{ *gpuDevice };
		if (!cmd) {
			return;
		}
		GPUCopyPass{ cmd }.UploadToGPUBuffer({ transferBuffer, 0 }, { *vertexBuffer, 0, sizeof(vertices) });
	}

	std::shared_ptr<GPUTexture> CreateImageTexture(const char* filename) {
		SDLSurface img{ IMG_Load(filename) };
		if (!img) {
			printf("load image[%s] failed! %s\n", filename, SDL_GetError());
			return {};
		}
		if (!img.ConvertToFormat(SDL_PIXELFORMAT_RGBA32)) {
			printf("convert image format failed! %s\n", SDL_GetError());
			return {};
		}

		GPUTransferBuffer transferBuffer{ *gpuDevice, SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD, (Uint32)(img.Height() * img.Pitch()) };
		if (!transferBuffer) {
			printf("Create transfer buffer failed! %s\n", SDL_GetError());
			return {};
		}
		memcpy(GPUTransferBufferMapper{ transferBuffer }, img.GetPixels(), (size_t)(img.Height() * img.Pitch()));

		GPUTextureCreateInfo texCreateInfo{};
		texCreateInfo.type = SDL_GPU_TEXTURETYPE_2D;
		texCreateInfo.format = SDL_GPU_TEXTUREFORMAT_R8G8B8A8_UNORM;
		texCreateInfo.usage = SDL_GPU_TEXTUREUSAGE_SAMPLER;
		texCreateInfo.width = (Uint32)img.Width();
		texCreateInfo.height = (Uint32)img.Height();
		texCreateInfo.layer_count_or_depth = 1;
		texCreateInfo.num_levels = 1;
		texCreateInfo.sample_count = SDL_GPU_SAMPLECOUNT_1;
		auto texture = std::make_shared<GPUTexture>(*gpuDevice, texCreateInfo);
		if (!*texture) {
			printf("Create texture failed! %s\n", SDL_GetError());
			return {};
		}

		GPUCommandBuffer cmd{ *gpuDevice };
		if (!cmd) {
			return {};
		}
		GPUTextureTransferInfo srcTexInfo{};
		srcTexInfo.transfer_buffer = transferBuffer;
		srcTexInfo.offset = 0;
		srcTexInfo.pixels_per_row = (Uint32)(img.Width());
		srcTexInfo.rows_per_layer = (Uint32)(img.Height());
		GPUTextureRegion dstTexRegion{};
		dstTexRegion.texture = *texture;
		dstTexRegion.mip_level = 0;
		dstTexRegion.layer = 0;
		dstTexRegion.x = 0;
		dstTexRegion.y = 0;
		dstTexRegion.z = 0;
		dstTexRegion.w = (Uint32)(img.Width());
		dstTexRegion.h = (Uint32)(img.Height());
		dstTexRegion.d = 1;
		GPUCopyPass{ cmd }.UploadToGPUTexture(srcTexInfo, dstTexRegion);
		return texture;
	}

	void CreateDepthTexture(Uint32 width, Uint32 height) {
		GPUTextureCreateInfo texCreateInfo{};
		texCreateInfo.type = SDL_GPU_TEXTURETYPE_2D;
		texCreateInfo.format = SDL_GPU_TEXTUREFORMAT_D16_UNORM;
		texCreateInfo.usage = SDL_GPU_TEXTUREUSAGE_DEPTH_STENCIL_TARGET;
		texCreateInfo.width = width;
		texCreateInfo.height = height;
		texCreateInfo.layer_count_or_depth = 1;
		texCreateInfo.num_levels = 1;
		texCreateInfo.sample_count = SDL_GPU_SAMPLECOUNT_1;
		depthTexture = std::make_unique<GPUTexture>(*gpuDevice, texCreateInfo);
		if (!*depthTexture) {
			printf("Create depth texture failed! %s\n", SDL_GetError());
			return;
		}
	}

	void InitPlanes() {
		{
			Plane plane;
			plane.color = glm::vec4(1, 1, 1, 1);
			plane.position = glm::vec3(0, 0, -0.5);
			plane.rotation = glm::vec3(-90, 0, 0);
			plane.scale = glm::vec3(10, 10, 10);
			plane.texture = floorTexture;

			planes.push_back(plane);
		}

		{
			Plane plane;
			plane.color = glm::vec4(0, 0, 0, 1);
			plane.position = glm::vec3(0.2, 0, -4);
			plane.rotation = glm::vec3(0, 0, 0);
			plane.texture = windowTexture;

			planes.push_back(plane);
		}

		{
			Plane plane;
			plane.color = glm::vec4(0.5, 0, 0, 1);
			plane.position = glm::vec3(-0.2, 0, -3);
			plane.rotation = glm::vec3(0, 0, 0);
			plane.texture = windowTexture;

			planes.push_back(plane);
		}
	}

	void Present() {
		camera.Update();
		mvp.view = camera.GetMat();

		if (window->GetWindowFlags() & SDL_WINDOW_MINIMIZED) {
			return;
		}
		GPUCommandBuffer cmd{ *gpuDevice };
		if (!cmd) {
			return;
		}
		GPUSwapchainTexture swapchainTexture{ cmd, *window };
		if (!swapchainTexture) {
			return;
		}
		GPUColorTargetInfo colorTargetInfo{};
		colorTargetInfo.texture = swapchainTexture;
		colorTargetInfo.clear_color = { 0.2f, 0.3f, 0.3f, 1.0f };
		colorTargetInfo.load_op = SDL_GPU_LOADOP_CLEAR;
		colorTargetInfo.store_op = SDL_GPU_STOREOP_STORE;
		colorTargetInfo.mip_level = 0;
		colorTargetInfo.layer_or_depth_plane = 0;
		colorTargetInfo.cycle = true;
		colorTargetInfo.cycle_resolve_texture = false;
		GPUDepthStencilTargetInfo depthStencilTargetInfo{};
		depthStencilTargetInfo.texture = *depthTexture;
		depthStencilTargetInfo.clear_depth = 1.0f;
		depthStencilTargetInfo.load_op = SDL_GPU_LOADOP_CLEAR;
		depthStencilTargetInfo.store_op = SDL_GPU_STOREOP_DONT_CARE;
		depthStencilTargetInfo.cycle = false;
		GPURenderPass renderPass{ cmd, &colorTargetInfo, 1, &depthStencilTargetInfo };
		if (!renderPass) {
			return;
		}

		renderPass.BindGraphicsPipeline(*gpuDevice->pipeline);
		renderPass.BindVertexBuffers(0, { *vertexBuffer, 0 });

		GPUViewport viewport{};
		viewport.x = 0.0f;
		viewport.y = 0.0f;
		viewport.w = (float)swapchainTexture.width;
		viewport.h = (float)swapchainTexture.height;
		viewport.min_depth = 0.0f;
		viewport.max_depth = 1.0f;
		renderPass.SetViewport(viewport);

		for (auto& plane : planes) {
			MVP local = mvp;

			local.model =
				glm::scale(
					glm::translate(
						glm::rotate(
							glm::rotate(
								glm::rotate(glm::mat4(1.0),
									glm::radians(plane.rotation.x), glm::vec3(1, 0, 0)),
								glm::radians(plane.rotation.y), glm::vec3(0, 1, 0)),
							glm::radians(plane.rotation.z), glm::vec3(0, 0, 1)),
						plane.position), plane.scale);
			renderPass.BindTextureSampler(0, { *plane.texture, *gpuDevice->texSampler });
			cmd.PushVertexUniformData(0, &local, sizeof(local));
			cmd.PushFragmentUniformData(0, &plane.color, sizeof(plane.color));
			renderPass.DrawPrimitives(6, 1, 0, 0);
		}
	}

	void Run() {
		const float speed = 0.1f;
		const float factor = 0.5f;

		sdlfx.onEvent = [this, speed, factor](EventData& event) {
			if (event.type == SDL_EVENT_KEY_DOWN) {
				if (event.key.key == SDLK_A) {
					camera.Move(glm::vec3(-speed, 0, 0));
				}
				if (event.key.key == SDLK_D) {
					camera.Move(glm::vec3(speed, 0, 0));
				}
				if (event.key.key == SDLK_W) {
					camera.Move(glm::vec3(-0, 0, -speed));
				}
				if (event.key.key == SDLK_S) {
					camera.Move(glm::vec3(0, 0, speed));
				}
				if (event.key.key == SDLK_ESCAPE) {
					event.type = SDL_EVENT_QUIT;
					SDL_PushEvent(&event);
				}
			}
			else if (event.type == SDL_EVENT_MOUSE_MOTION) {
				camera.RotateX(-event.motion.yrel * factor);
				camera.RotateY(-event.motion.xrel * factor);
			}
			return true;
		};
		sdlfx.onFrame = [this]() {
			Present();
		};

		sdlfx.Run();
	}

	void Cleanup() {
		gpuDevice->WaitIdle();
	}
};

int main() {
	try {
		AppContext app;
		if (app.Init()) {
			app.Run();
		}
		app.Cleanup();
	}
	catch (const std::exception& e) {
		printf("App error! '%s'\n", e.what());
		return -1;
	}
	return 0;
}